import json
import telnetlib
import time
import re
from datetime import datetime
import easygui
import requests


class Mail:
    def __init__(self,app_token,uid_myself):
        self.app_token = app_token  # 本处改成自己的应用 self.app_token
        self.uid_myself = uid_myself  # 本处改成自己的 UID

    def wxpusher_send_by_webapi(self, msg):
        """利用 wxpusher 的 web api 发送 json 数据包，实现微信信息的发送"""
        webapi = 'http://wxpusher.zjiecode.com/api/send/message'
        data = {
            "appToken": self.app_token,
            "content": msg,
            "summary": msg[:10],  # 该参数可选，默认为 msg 的前10个字符
            "contentType": 1,
            "uids": [self.uid_myself],
        }
        result = requests.post(url=webapi, json=data)
        return result.text

def to_date_str(time_str):
    time_format = "%b %d %Y %H:%M:%S%z"
    dt = datetime.strptime(time_str, time_format)
    target_format = "%Y-%m-%d %H:%M:%S"
    formatted_dt = dt.strftime(target_format)
    return formatted_dt


def get_new_threat(tn,model='USG6500'):
    for threat_key,threat_type in threat_types.items():
        tn.write(threat_type.encode('utf-8') + b'\n')
        time.sleep(2)
        res = tn.read_very_eager().decode('GBK')
        lens=re.findall('(.*?)\n', res)
        for len in lens:
            Policy = (re.findall('Policy="(.*?)",', len))
            if Policy:
                start_time = re.findall(f'(.*?) {model}', len)
                by_time = to_date_str(start_time[0])
                # 进入策略名称
                Policy_type =  Policy[0]
                # 对方IP地址
                SrcIp = (re.findall('SrcIp=(.*?),', len))[0]
                # 内网服务器地址
                DstIp = (re.findall('DstIp=(.*?),', len))[0]
                # 对方端口
                SrcPort = (re.findall('SrcPort=(.*?),', len))[0]
                # 内网服务器端口
                DstPort = (re.findall('DstPort=(.*?),', len))[0]
                # 来源区域
                SrcZone = (re.findall('SrcZone=(.*?),', len))[0]
                # 内网区域
                DstZone = (re.findall('DstZone=(.*?),', len))[0]
                # 入侵方式
                SignName = (re.findall('SignName="(.*?)",', len))[0]
                yield {
                    '入侵方式':threat_key,
                    '手段': SignName,
                    'time':by_time,
                    '进入策略名称':Policy_type,
                    '来源区域':SrcZone,
                    '对方IP地址':SrcIp,
                    '对方端口':SrcPort,
                    '内网区域':DstZone,
                    '内网服务器地址':DstIp,
                    '内网服务器端口':DstPort,
                }

def connect_tn(user,pwd,ip='172.24.121.254',t=1,):
    """
    链接到数据库
    执行tw需要时间
    :return:
    """
    # 创建连接对象，链接IP地址
    tn = telnetlib.Telnet(ip)
    tn.write(user.encode('utf-8') + b'\n')
    time.sleep(t)
    tn.write(pwd.encode('utf-8') + b'\n')
    time.sleep(t)
    return tn


def send_f(data_ary,app_token,uid_myself):
    send_text= f'新威胁次数：{len(data_ary)}\n'
    for new_threat_lien_k, new_threat_lien_v in data_ary[0].items():
        send_text += f"{new_threat_lien_k}:{new_threat_lien_v}\n"
    print(send_text)
    wx=Mail(app_token,uid_myself)
    wx.wxpusher_send_by_webapi(msg=send_text)


if __name__ == '__main__':
    app_token = ''  # 本处改成自己的应用百度wxpusher
    uid_myself = ''  # 本处改成自己的 UID 百度wxpusher
    ip = ''  # 设备ip
    user = ''  # 用户
    pwd = easygui.enterbox("密码:")
    mod = ''  # 设备型号
    stime = 1200 # 暂停获取时间
    while True:
        # 读取JSON文件
        with open('data.json', 'r') as f:
            data = json.load(f)
        # 查询语句
        threat_types={'入侵':'display logbuffer sec-log | include An intrusion was detected',
                      '病毒':'display logbuffer sec-log | include A virus was detected',  # 因为厦门很少遇到病毒 virus仅仅是猜测
                      '僵木蠕':'display logbuffer sec-log | include A botnet was detected'}
        # 链接设备
        tn = connect_tn(user=user,pwd=pwd,ip=ip)
        # 获取最新的异常
        new_threats = get_new_threat(tn,model=mod)
        add_data_ary = [] # 添加到json中的新数据
        for new_threat in new_threats:
            if new_threat['time'] not in data:
                data[new_threat['time']]=new_threat
                add_data_ary.append(new_threat)
            else:
                break
        # 保存修改后的数据回到 JSON 文件
        with open('data.json', 'w') as f:
            json.dump(data, f, indent=4)
        # 发送add_data_ary
        if len(add_data_ary)>0:
            send_f(add_data_ary,app_token,uid_myself)
        print(f"等待下次执行时间{stime}秒后....")
        time.sleep(stime)
